#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/param.h>
#include <sys/time.h>
#include <arpa/inet.h>
#include <netdb.h>

#include "config.h"
#include "setup_ib.h"
#include "helper.h"

struct IBRes ib_res;
char *server_name;

struct IBDest local_dest;
struct IBDest remote_dest;

int get_server_dest()
{
    int sockfd;
    struct addrinfo hints, *result, *rp;

    memset(&hints, 0, sizeof(struct addrinfo));
    hints.ai_family = AF_UNSPEC; // Allow IPv4 or IPv6.
    hints.ai_socktype = SOCK_STREAM; // TCP

    if (getaddrinfo(server_name, PORT_STR, &hints, &result)) {
        perror("Found address failed!\n");
        return 1;
    }

    // Start connection:
    for (rp = result; rp; rp = rp->ai_next) {
        sockfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if (sockfd == -1)
            continue;
        if (!connect(sockfd, rp->ai_addr, rp->ai_addrlen))
            break; // Success.
        close(sockfd);
    }

    if (!rp) {
        perror("Connection with the server failed.\n");
        return 1;
    }

    freeaddrinfo(result);

    // Get server details:
    {
        char *buf;
        buf = (char *)malloc(sizeof(struct IBDest));
        int offset = 0;

        // Send my details:
        memcpy(buf, &local_dest, sizeof(struct IBDest));
        write(sockfd, buf, sizeof(struct IBDest));

        // Get him details:
        memset(buf, 0, sizeof(struct IBDest));
        offset = 0;

        while (offset < sizeof(struct IBDest))
            offset += read(sockfd, buf + offset, sizeof(struct IBDest) - offset);

        memcpy(&remote_dest, buf, sizeof(struct IBDest));
    }
    // Finish connection:
    close(sockfd);
    return 0;
}

int get_client_dest()
{
    int sockfd, connfd, len;
    struct sockaddr_in server_address, client_address;
    // Start connection:
    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd == -1)
    {
        perror("Socket creation failed!\n");
        return 1;
    }

    memset(&server_address, 0, sizeof(server_address));
    server_address.sin_family = AF_INET;
    server_address.sin_addr.s_addr = htonl(INADDR_ANY);
    server_address.sin_port = htons(PORT);

    if (bind(sockfd, (struct sockaddr*)&server_address, sizeof(server_address))) {
        perror("Socket bind failed!\n");
        return 1;
    }

    if (listen(sockfd, 5)) {
        perror("Listen failed...\n");
        return 1;
    }

    len = sizeof(client_address);
    connfd = accept(sockfd, (struct sockaddr*)&client_address, &len);
    if (connfd < 0) {
        perror("Server accept failed!\n");
        return 1;
    }

    // Get client details:
    {
        char *buf;
        buf = (char*)malloc(sizeof(struct IBDest));
        int offset = 0;

        // Get him details.
        while (offset < sizeof(struct IBDest))
            offset += read(connfd, buf + offset, sizeof(struct IBDest) - offset);

        memcpy(&remote_dest, buf, sizeof(struct IBDest));
        // Send my details:
        memset(buf, 0, sizeof(struct IBDest));
        offset = 0;

        memcpy(buf, &local_dest, sizeof(struct IBDest));
        write(connfd, buf, sizeof(struct IBDest));
    }

    // Finish connection:
    close(sockfd);
    return 0;
}

int connect_between_qps()
{
    struct ibv_qp_attr qp_attr = {
            .qp_state       = IBV_QPS_RTR,
            .path_mtu       = MTU,
            .dest_qp_num    = remote_dest.qpn,
            .rq_psn         = remote_dest.psn,
            .max_dest_rd_atomic = 1,
            .min_rnr_timer      = 12,
            .ah_attr        = {
                    .is_global  = 0,
                    .dlid       = remote_dest.lid,
                    .sl     = 0,
                    .src_path_bits  = 0,
                    .port_num   = IB_PORT
            }
    };

    if (remote_dest.gid.global.interface_id) {
        qp_attr.ah_attr.is_global = 1;
        qp_attr.ah_attr.grh.hop_limit = 1;
        qp_attr.ah_attr.grh.dgid = remote_dest.gid;
        qp_attr.ah_attr.grh.sgid_index = 1;
    }

    if (visroce_modify_qp(ib_res.qp, &qp_attr, IBV_QP_STATE | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | IBV_QP_RQ_PSN |
                                           IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER | IBV_QP_AV)) {
        perror("Failed to modify QP to RTR.\n");
        return 1;
    }

    qp_attr.qp_state        = IBV_QPS_RTS;
    qp_attr.timeout     = 14;
    qp_attr.retry_cnt       = 7;
    qp_attr.rnr_retry       = 7;
    qp_attr.sq_psn          = local_dest.psn;
    qp_attr.max_rd_atomic  = 1;
    if (visroce_modify_qp(ib_res.qp, &qp_attr, IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT |
                                           IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC)) {
        perror("Failed to modify QP to RTS.\n");
        return 1;
    }
    return 0;
}

int setup_ib()
{
    struct ibv_device **dev_list, *ib_dev;
    union ibv_gid temp_gid;
    char             gid[33];
    srand48(getpid() * time(NULL));
    memset(&ib_res, 0, sizeof(struct IBRes));
    ib_res.size = BUFFER_SIZE;
    ib_res.buf = malloc(roundup(BUFFER_SIZE, sysconf(_SC_PAGESIZE)));
    if (!ib_res.buf) {
        perror("Couldn't allocate work buf.\n");
        return 1;
    }

    memset(ib_res.buf, 0x7b + (server_name ? 0 : 1), BUFFER_SIZE);
    // Get an IB devices list:
    dev_list = visroce_get_device_list(NULL);
    if (!dev_list) {
        perror("Failed to get IB devices list.\n");
        return 1;
    }

    // Get an IB device:
    ib_dev = *dev_list;
    if (!ib_dev) {
        perror("No IB devices found.\n");
        return 1;
    }
    printf("ib device name: %s\n", visroce_get_device_name(ib_dev));

    // Open an IB device context:
    ib_res.ctx = visroce_open_device(*dev_list);
    if (!ib_res.ctx) {
        fprintf(stderr, "Couldn't get context for %s.\n", visroce_get_device_name(ib_dev));
        return 1;
    }

    // Allocate a Protection Domain:
    ib_res.pd = visroce_alloc_pd(ib_res.ctx);
    if (!ib_res.pd) {
        perror("Failed to allocate Protection Domain.\n");
        return 1;
    }

    // Query IB port attribute:
    if (visroce_query_port(ib_res.ctx, IB_PORT, &ib_res.port_info)) {
        perror("Failed to query IB port information.\n");
        return 1;
    }

    // Register a Memory Region:
    ib_res.mr = visroce_reg_mr(ib_res.pd, ib_res.buf, ib_res.size,
                        IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE);
    if (!ib_res.mr) {
        perror("Couldn't register Memory Region.\n");
        return 1;
    }

    printf("rkey: 0x%06x, lkey: 0x%06x\n",  ib_res.mr->rkey, ib_res.mr->lkey);

    // Creates a Completion Queue:
    ib_res.cq = visroce_create_cq(ib_res.ctx, 2 * COUNT, NULL, NULL, 0);
    if (!ib_res.cq) {
        perror("Couldn't create Completion Queue.\n");
        return 1;
    }

    // Creates a Queue Pair:
    {
        struct ibv_qp_attr attr;
        struct ibv_qp_init_attr qp_init_attr = {
                .send_cq = ib_res.cq,
                .recv_cq = ib_res.cq,
                .cap = {
                        .max_send_wr = COUNT,
                        .max_recv_wr = COUNT,
                        .max_send_sge = 1,
                        .max_recv_sge = 1,
                },
                .qp_type = IBV_QPT_RC,
        };

        ib_res.qp = visroce_create_qp(ib_res.pd, &qp_init_attr);
        if (!ib_res.qp) {
            perror("Couldn't create Queue Pair.\n");
            return 1;
        }

        struct ibv_qp_attr qp_attr = {
                .qp_state        = IBV_QPS_INIT,
                .pkey_index      = 0,
                .port_num        = IB_PORT,
                .qp_access_flags = IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE
        };

        if (visroce_modify_qp(ib_res.qp, &qp_attr, IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS)) {
            perror("Failed to modify QP to INIT.\n");
            return 1;
        }
    }

    // Query IB port attribute:
    if (visroce_query_port(ib_res.ctx, IB_PORT, &ib_res.port_info)) {
        perror("Failed to query IB port information.\n");
        return 1;
    }
    printf("ib device port info, active mtu: %d\n", ib_res.port_info.active_mtu);

    local_dest.lid = ib_res.port_info.lid;
    if (ib_res.port_info.link_layer != IBV_LINK_LAYER_ETHERNET && !local_dest.lid) {
        perror("Couldn't get LID.\n");
        return 1;
    }

    // Get LID:
    if (visroce_query_gid(ib_res.ctx, IB_PORT, 1, &temp_gid)) {
        printf("ibv query gid failed\n");   
        return 1;
    }
    memcpy(local_dest.gid.raw, temp_gid.raw ,16);

    printf("gid : %02x:%02X:%d:%d:%d:%d\n", temp_gid.raw[10], temp_gid.raw[11], temp_gid.raw[12], temp_gid.raw[13], temp_gid.raw[14], temp_gid.raw[15]);

    // Get QPN:
    local_dest.qpn = ib_res.qp->qp_num;

    // Set PSN:
    local_dest.psn = lrand48() & 0xffffff;
    // generate gid
    local_dest.gid_index = 1;

    // Get details of each side:
    if (server_name) {
        if (get_server_dest()) {
            perror("Couldn't get server's details.\n");
            return 1;
        }
    }
    else{
        if (get_client_dest()) {
            perror("Couldn't get client's details.\n");
            return 1;
        }
    }
    inet_ntop(AF_INET6, &local_dest.gid, gid, sizeof(gid));
    printf("local dest, qpn: %d, psn: %d, lid: %d, gid: %s\n", local_dest.qpn, local_dest.psn, local_dest.lid,gid);
    inet_ntop(AF_INET6, &remote_dest.gid, gid, sizeof(gid));
    printf("remote dest, qpn: %d, psn: %d, lid: %d, gid: %s\n", remote_dest.qpn, remote_dest.psn, remote_dest.lid, gid);

    // Connect Queue Pair:
    if (connect_between_qps()) {
        perror("Couldn't connect between client and server.\n");
        return 1;
    }

    visroce_free_device_list(dev_list);
    return 0;
}

void close_ib_connection()
{
    if (ib_res.qp){
        visroce_destroy_qp(ib_res.qp);
    }

    if (ib_res.cq){
        visroce_destroy_cq(ib_res.cq);
    }

    if (ib_res.mr) {
        visroce_dereg_mr(ib_res.mr);
    }

    if (ib_res.pd) {
        visroce_dealloc_pd(ib_res.pd);
    }

    if (ib_res.ctx) {
        visroce_close_device(ib_res.ctx);
    }

    if (ib_res.buf){
        free(ib_res.buf);
    }
}
